/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.br.comps;

import coalescent.br.AMCElement;
import coalescent.br.BRWCompInfo;

import commons.util.FormatUtil;

import java.math.BigDecimal;
import java.util.HashMap;
import java.util.Map;

/**
 * Computes the probability of a sample configuration by observing its recursion (via the
 * post-order implementation).
 *
 * @author Susanta Tewari
 * @version 1.0 Dec 3, 2010
 *
 * @param <E>
 * @param <R>
 */
public final class BRWProbComputer<E extends Enum<E>, R extends AMCElement<E, R>>
        extends BRWCompAdapter<E, R> {

    /**
     * It is required when adding the probabilities of ancestral configurations for a given
     * event type
     */
    private Map<R, BigDecimal> adder = new HashMap<>();

    /**
     * cache of counts of genealogies for the ancestral configurations
     */
    private Map<R, BigDecimal> genealogyProbCache = new HashMap<>();

    /** Field description */
    private BigDecimal probability;

    /**
     * The probability of the sample configuration being observed.
     *
     * @return probability of the sample configuration being observed.
     */
    public BigDecimal probability() {
        return probability;
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onBoundaryCondn(BRWCompInfo<E, R> event) {

        R sampleConfig = event.getCurrentConfig();

        if (genealogyProbCache.get(sampleConfig) == null) {
            genealogyProbCache.put(sampleConfig, sampleConfig.sinkProbability());
        }
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onStartingIteratingEvents(BRWCompInfo<E, R> event) {


        // check if the sample config is in cache
        R sampleConfig = event.getCurrentConfig();


        // initialize the prob
        genealogyProbCache.put(sampleConfig, BigDecimal.valueOf(0.0));
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onIteratingSingleEventTypes(BRWCompInfo<E, R> event) {

        R sampleConfig = event.getCurrentConfig();

        adder.put(sampleConfig, BigDecimal.valueOf(0.0));
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onIteratingAncestralConfigs(BRWCompInfo<E, R> event) {

        R sampleConfig           = event.getCurrentConfig();
        R ancestralConfig        = event.getAncestralConfig();
        BigDecimal sampleProb = sampleConfig.transitionProb(event.getEventType(),
                                    event.getAncestralConfig());
        BigDecimal ancestralProb = genealogyProbCache.get(ancestralConfig);
        BigDecimal value         = sampleProb.multiply(ancestralProb);
        BigDecimal previous      = genealogyProbCache.get(sampleConfig);

        genealogyProbCache.put(sampleConfig, previous.add(value));
    }

    /**
     * Method description
     *
     *
     * @param event
     */
    @Override
    public void onFinishedIteration(BRWCompInfo<E, R> event) {

        super.onFinishedIteration(event);

        R sampleConfig = event.getCurrentConfig();

        if (event.getSampleConfig().equals(sampleConfig)) {

            probability = genealogyProbCache.get(sampleConfig);


            // de-allocate the resources
            genealogyProbCache = null;
        }
    }

    /**
     * Method description
     *
     *
     * @return
     */
    @Override
    public String toString() {

        super.toString();

        return "Exact Probability: " + FormatUtil.format(probability);
    }
}
